import concurrent.futures
import datetime
import functools
import json
import logging
import os
import traceback
import typing

from absl import (app, flags)
import pandas as pd
import pytz

import coin.proto.coin_data_replay_pb2 as cdr_pb2
from coin.base.logging import init_logging
from coin.base.param_util import to_list
from coin.base.query_util import query_exchange_rates
from coin.base.datetime_util import (
    convert_string_to_datetime,
    get_dt_ranges,
    iterate_date,
    to_timestamp_int,)
from coin.proto.coin_strategy_summary_pb2 import (
    StrategySummary, StrategySummaryMessage)
from coin.support.accounting.logic.constants import COIN_LIST
from coin.support.data_replay.driver import DataReplayDriver
from coin.support.pta.logic.pta_dumper import (PtaDumpStrategy,
                                               aggregate_strategy_detail,
                                               dump_pta_stats,)
from coin.support.pta.util.info_util import get_strat_info_list
from coin.tool.strat_monitor.util.calculator_util import adjust_pnl_from_config
from coin.tool.strat_monitor.util.monitor_util import (
    get_active_strategy,
    get_strat_from_json_config,
    get_pnl_adjust_info,
    filter_empty_summary)
from coin.util.queue.config import KafkaConfig
from coin2.service.strategy.trading_summary_service_pb2 import (
    QueryTradingSummaryIntervalHistoriesRequestProto)
from xunkemgmt_client.client.api_client import XunkemgmtClient
from xunkemgmt_client.client.util.query_util import (
    query_accounts,
    query_exchange_apis,
    query_strategies)
from xunkemgmt_client.support.accounting.logic.constants import (
    FIAT_LIST, STABLE_COIN_LIST)

FLAGS = flags.FLAGS


def _check_no_price(pta_summary):
  assert isinstance(pta_summary, StrategySummary)
  for _, stat in pta_summary.stat_per_quote.items():
    if stat.pnl_mark != 0:
      assert stat.HasField('pnl_mark_in_usd'), pta_summary
    if stat.turnover_mark != 0:
      assert stat.HasField('turnover_mark_in_usd'), pta_summary


def _get_source_info(strat_log_root_dir, kafka_config):
  source_info = strat_log_root_dir if strat_log_root_dir \
                is not None else ','.join(kafka_config.kafka_servers)
  return source_info


def _get_strat_log_size(root_dir, strat_info):
  total_size = 0
  log_dir = os.path.join(root_dir, strat_info.machine, strat_info.trading_date)
  for each_log in os.listdir(log_dir):
    if each_log.startswith('%s.' % strat_info.strategy_name):
      total_size += os.path.getsize(os.path.join(log_dir, each_log))
  return total_size


def launch_replay_func(*, driver_config, pta_strategy_maker):
  pta_strategy = pta_strategy_maker()
  assert len(driver_config.requests) == 1
  data_request = driver_config.requests[0]
  strat_info = pta_strategy.strat_info
  driver = DataReplayDriver(driver_config, pta_strategy)
  driver.run()
  pta_summaries = pta_strategy.pta_summaries
  if data_request.HasField('root_dir'):
    log_size_byte = _get_strat_log_size(data_request.root_dir, strat_info)
    for pta_summary in pta_summaries:
      pta_summary.source.strategy_log_info.log_size_byte = log_size_byte
  return pta_summaries, pta_strategy.pta_details


def calculate_pta_stats_from_strat_log(*,
                                       start_dt,
                                       end_dt,
                                       strategy_list,
                                       topic_type,
                                       kafka_config_filename,
                                       strat_log_root_dir,
                                       max_workers,
                                       interval_sec=None,
                                       offset_hour=None,
                                       ignore_error=False,
                                       pnl_adj=None,
                                       realtime_pnl_adj=False,
                                       account_universe=None,
                                       exchange_api_universe=None,
                                       strategy_universe=None,
                                       ticker=None,
                                       plot_dir=None,
                                       plot_order=False,
                                       log_history_dir=None,
                                       strat_stat_log_dir=None):
  func_list = []
  kafka_config = None
  if kafka_config_filename is not None:
    kafka_config = KafkaConfig.from_cmd_config(kafka_config_filename)
  strat_infos = get_strat_info_list(strat_log_root_dir=strat_log_root_dir,
                                    kafka_config=kafka_config,
                                    strategy_list=strategy_list,
                                    start_date=start_dt.date(),
                                    end_date=end_dt.date())
  source_info = _get_source_info(strat_log_root_dir, kafka_config)
  offset_hour = offset_hour if offset_hour is not None else 3
  save_order_event = True if log_history_dir is not None else False
  for strat_info in strat_infos:
    strategy_maker = functools.partial(PtaDumpStrategy,
                                       strat_info=strat_info,
                                       start_time=start_dt,
                                       end_time=end_dt,
                                       source_info=source_info,
                                       ignore_error=ignore_error,
                                       pnl_adj=pnl_adj,
                                       realtime_pnl_adj=realtime_pnl_adj,
                                       account_universe=account_universe,
                                       exchange_api_universe=exchange_api_universe,
                                       strategy_universe=strategy_universe,
                                       ticker=ticker,
                                       save_order_event=save_order_event,
                                       plot_dir=plot_dir,
                                       plot_order=plot_order,
                                       log_history_dir=log_history_dir,
                                       strat_stat_log_dir=strat_stat_log_dir)
    root_dir = strat_log_root_dir if kafka_config_filename is None else None
    data_req = cdr_pb2.DataReplayRequest(
        topic_type=topic_type,
        root_dir=root_dir,
        kafka_config_file=kafka_config_filename,
        strategy_names=to_list(strat_info.strategy_name),
        machines=to_list(strat_info.machine))
    driver_config = cdr_pb2.DataReplayDriverConfig(
        start_timestamp=to_timestamp_int(start_dt - datetime.timedelta(hours=offset_hour)),
        end_timestamp=to_timestamp_int(end_dt),
        interval_sec=interval_sec,
        requests=to_list(data_req))
    replay_func = functools.partial(launch_replay_func,
                                    driver_config=driver_config,
                                    pta_strategy_maker=strategy_maker)
    func_list.append((strat_info, replay_func))

  num_func = len(func_list)
  pta_summaries = []
  pta_details = []
  success = 0
  if max_workers > 1:
    with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
      fut_to_strat_info = {executor.submit(func): key for key, func in func_list}
      done_futs, not_done_futs = concurrent.futures.wait(fut_to_strat_info.keys(), timeout=None)
      func_list = [(fut_to_strat_info[fut], fut.result) for fut in done_futs]
      for fut in not_done_futs:
        strat_info = fut_to_strat_info[fut]
        logging.error('Fail to dump strategy log due to timeout. %s', strat_info._asdict())
      for pid, process in executor._processes.items():
        process.terminate()

  for key, func in func_list:
    try:
      summaries, details = func()
      pta_summaries.extend(summaries)
      pta_details.extend(details)
      success += 1
    except Exception as e:
      logging.error('Fail to dump strategy log. %s, %s, %s', key._asdict(), type(e), e)
      logging.error(traceback.format_exc())
  logging.info('Success rate: %s/%s' % (success, num_func))
  num_fail = num_func - success
  return pta_summaries, pta_details, num_fail


def query_ticker(base_list: typing.List[str],
                 quote: str,
                 trading_date: datetime.date) -> typing.Dict[str, float]:
  price_list = query_exchange_rates(base_list, quote, trading_date)
  '''
  na_price_ccys = [base for base, price in zip(base_list, price_list) if price is None]
  assert len(na_price_ccys) == 0, na_price_ccys
  '''
  ticker = {'%s-%s' % (base, quote): price for base, price in zip(base_list, price_list)}
  return ticker


def gen_strategy_summary_message(trading_date, result_list):
  td = datetime.datetime.combine(trading_date, datetime.time.max)
  ts = int(td.replace(tzinfo=pytz.UTC).timestamp() * 10**9)
  return StrategySummaryMessage(fetched_ts=ts, summaries=result_list)


def main(_):
  init_logging()
  start_time = FLAGS.start_time
  assert start_time, '--start_time must be specified.'
  end_time = FLAGS.end_time
  assert end_time, '--end_time must be specified.'
  assert FLAGS.business_unit, '--business_unit must be specified.'
  pta_db_config = FLAGS.pta_db_config
  latency_db_config = FLAGS.latency_db_config
  strategy_config_filename = FLAGS.strategy_config_filename
  kafka_config_filename = FLAGS.kafka_config_filename

  account_universe = query_accounts(as_proto=True)
  exchange_api_universe = query_exchange_apis(as_proto=True)
  strategy_universe = query_strategies(as_proto=True)

  start_dt = convert_string_to_datetime(start_time)
  end_dt = convert_string_to_datetime(end_time)
  dt_ranges = get_dt_ranges(start_dt, end_dt, FLAGS.aggregate)

  pnl_adj = get_pnl_adjust_info(
      FLAGS.pnl_adj_file, start_dt.date(), end_dt.date() + datetime.timedelta(days=1))

  if strategy_config_filename is not None:
    strat_list = get_strat_from_json_config(strategy_config_filename)
  else:
    threshold = to_timestamp_int(start_dt)
    strat_list = get_active_strategy(update_timestamp_threshold=threshold)

  filtered = strategy_universe
  business_units = [elem.strip() for elem in FLAGS.business_unit.split(',')]
  filtered = [elem for elem in filtered if elem.business_unit in business_units]
  if FLAGS.market_type is not None:
    market_types = [elem.strip() for elem in FLAGS.market_type.split(',')]
    with XunkemgmtClient() as xunkemgmt_client:
      rsp = xunkemgmt_client.query_trading_summary_interval_histories(
          QueryTradingSummaryIntervalHistoriesRequestProto(
              business_units=business_units,
              market_types=market_types,
              start_timestamp=to_timestamp_int(start_dt),
              end_timestamp=to_timestamp_int(end_dt),
              summary_info_only=True,
              agg_types=['STRATEGY_NAME']))
    strategy_names = [history.summary_info.strategy_name for history in rsp.histories]
    filtered = [elem for elem in filtered if elem.strategy_name in strategy_names]
  if FLAGS.exchange is not None:
    exchanges = [elem.strip() for elem in FLAGS.exchange.split(',')]
    with XunkemgmtClient() as xunkemgmt_client:
      rsp = xunkemgmt_client.query_trading_summary_interval_histories(
          QueryTradingSummaryIntervalHistoriesRequestProto(
              business_units=business_units,
              exchanges=exchanges,
              start_timestamp=to_timestamp_int(start_dt),
              end_timestamp=to_timestamp_int(end_dt),
              summary_info_only=True,
              agg_types=['STRATEGY_NAME']))
    strategy_names = [history.summary_info.strategy_name for history in rsp.histories]
    filtered = [elem for elem in filtered if elem.strategy_name in strategy_names]
  if FLAGS.strategy_group is not None:
    strategy_groups = [elem.strip() for elem in FLAGS.strategy_group.split(',')]
    filtered = [elem for elem in filtered if elem.strategy_group in strategy_groups]
  strat_list = set(strat_list).intersection(
      set([elem.strategy_name for elem in filtered]))
  if FLAGS.strategy_name is not None:
    strat_list = [elem.strip() for elem in FLAGS.strategy_name.split(',')]

  total_num_fail = 0
  for start_dt, end_dt in dt_ranges:
    start_td_str = start_dt.strftime('%Y%m%dT%H%M%S')
    end_td_str = end_dt.strftime('%Y%m%dT%H%M%S')
    logging.info('\nRunning for %s-%s ...' % (start_td_str, end_td_str))
    base_list = list(set(COIN_LIST + FIAT_LIST + STABLE_COIN_LIST))
    ticker = query_ticker(base_list, 'USD', trading_date=end_dt.date())

    pta_summaries, details, num_fail = calculate_pta_stats_from_strat_log(
        start_dt=start_dt,
        end_dt=end_dt,
        strategy_list=strat_list,
        topic_type=cdr_pb2.TopicType.Value(FLAGS.topic_type),
        kafka_config_filename=kafka_config_filename,
        strat_log_root_dir=FLAGS.root_dir,
        interval_sec=FLAGS.interval_sec,
        offset_hour=FLAGS.offset_hour,
        max_workers=FLAGS.max_workers,
        ignore_error=FLAGS.ignore_error,
        pnl_adj=pnl_adj,
        realtime_pnl_adj=FLAGS.realtime_pnl_adj,
        account_universe=account_universe,
        exchange_api_universe=exchange_api_universe,
        strategy_universe=strategy_universe,
        ticker=ticker,
        plot_dir=FLAGS.plot_dir,
        plot_order=FLAGS.plot_order,
        log_history_dir=FLAGS.log_history_dir,
        strat_stat_log_dir=FLAGS.strat_stat_log_dir)
    total_num_fail += num_fail
    pta_summaries = filter_empty_summary(pta_summaries)
    if len(pta_summaries) == 0:
      continue
    if FLAGS.check_no_price:
      for pta_summary in pta_summaries:
        _check_no_price(pta_summary)
    hourly_details = []
    hourly_details = [
        aggregate_strategy_detail(detail, 60) for detail in details if len(detail.each_detail) > 0
    ]
    if not FLAGS.realtime_pnl_adj:
      pta_summaries, hourly_details = adjust_pnl_from_config(
          pta_summaries, hourly_details, start_dt.date(), end_dt.date(), pnl_adj)
    summary_msg = gen_strategy_summary_message(end_dt.date(), pta_summaries)
    dump_pta_stats(summary_msg,
                   hourly_details,
                   ticker=ticker,
                   output_file=FLAGS.output_file,
                   output_csv_file=FLAGS.output_csv_file,
                   slack_receiver=FLAGS.slack_receiver,
                   pta_db_config_path=pta_db_config,
                   latency_db_config_path=latency_db_config)
  return total_num_fail


if __name__ == '__main__':
  flags.DEFINE_string('start_time', None, 'Start date in form of %Y%m%d or %Y%m%dT%H%M%S.')

  flags.DEFINE_string('end_time', None, 'End date in form of %Y%m%d or %Y%m%dT%H%M%S.')

  flags.DEFINE_integer('max_workers', 24, 'Max number of workers.')

  flags.DEFINE_boolean('ignore_error', False, 'Ignore error')

  flags.DEFINE_float('offset_hour',
                     None,
                     'offset hour between the start time to read data and start time of pta.')

  flags.DEFINE_string('strategy_name', None, 'e.g. lm_agg2_binance_spot_minor_4,basis_smm2_prex_sandbox_binance_spot')

  flags.DEFINE_string('strategy_group', None, 'filter by strategy_group')

  flags.DEFINE_string('market_type', None, 'filter by market_type')

  flags.DEFINE_string('exchange', None, 'filter by exchange')

  flags.DEFINE_string('business_unit', None, 'filter by business_unit')

  flags.DEFINE_boolean('aggregate', False,
                       'calculate total pnl from start_date to end_date.')

  flags.DEFINE_boolean('check_no_price', False,
                       'raise exception when price in usd is NA.')

  flags.DEFINE_integer('interval_sec', None, 'interval_sec')

  flags.DEFINE_string('root_dir', None, '')

  flags.DEFINE_string('kafka_config_filename', None, 'kafka config')

  flags.DEFINE_string('topic_type', 'STRATEGY', '')

  flags.DEFINE_string('pnl_adj_file', None, 'pnl_adj_file')

  flags.DEFINE_boolean('realtime_pnl_adj', False, 'realtime pnl adj')

  flags.DEFINE_string('pta_db_config', None, '')

  flags.DEFINE_string('latency_db_config', None, 'path of latency db config.')

  flags.DEFINE_string('output_file', None, 'output_file')

  flags.DEFINE_string('output_csv_file', None, '')

  flags.DEFINE_string('slack_receiver', None, 'slack recevier')

  flags.DEFINE_string('log_history_dir', None, 'log_history_dir')

  flags.DEFINE_string('strat_stat_log_dir', None, 'strat_stat_log_dir')

  flags.DEFINE_string('plot_dir', None, 'plot dir')

  flags.DEFINE_boolean('plot_order', False, 'do order plot')

  flags.DEFINE_string('strategy_config_filename',
                      None,
                      'strategy_config_filename')

  flags.DEFINE_string('kline_db_config',
                      '../../coin_deploy/support_info/db_config_iosg/mysql_config_prod.json',
                      '')

  flags.DEFINE_string('kline_symbols_list',
                      '../../coin_deploy/support_info/kline_config/kline_symbols_list.csv',
                      'kline_symbols_list')

  app.run(main)
